package perf;

import editor.ErrorHandler;
import editor.PwnIDE;
import lang.Encoder;
import lang.asm.ASM;
import model.Project;
import model.code.CodeFileChunk;
import model.code.Function;
import model.code.GlobalVariable;
import model.code.Section;
import model.code.line.*;
import model.symbol.LocalVariable;
import model.symbol.Scope;
import util.ArrayUtil;
import util.FileUtil;

import java.io.File;
import java.io.IOException;
import java.text.DecimalFormat;
import java.util.ArrayList;
import java.util.Calendar;
import java.util.GregorianCalendar;

/**
 * This class is dedicated to taking a performance test Function and building a test executable for it.
 * It specifies special Functions for use by performance tests and to be the entry point that calls a performance test.
 * @author Neil Dickson
 */
public class PerformanceTestEncoder {
	private final Function			rootFunction;
	public final Function			startFunction;
	public final Function			stopFunction;
	private final GlobalVariable	timesArrayVar;
	private final GlobalVariable	filenameVar;
	private final GlobalVariable	timeVar;
	private final GlobalVariable	raxVar;
	private final GlobalVariable	rdxVar;
	private final GlobalVariable	xmm0Var;
	private final GlobalVariable	addressVar;

	private static final DecimalFormat	YEAR_FORMAT = new DecimalFormat("0000");
	private static final DecimalFormat	NUMBER_FORMAT = new DecimalFormat("00");

	public static final int	MAX_DATA_LINE_NUM		= 2;
	public static final int	TEST_FUNCTION_LINE_NUM	= 7;
	private static final String[]	rootFunctionText = {
			"sub\trsp,28h",
			"xor\trcx,rcx",
			"mov\trdx,",	// max num data goes here
			"shl\trdx,4",
			"call\tLocalAlloc",
			"add\trsp,28h",
			"mov\t__pPTestTimesArray,rax",
			"call\t",		// test function name goes here
			"mov\trbx,rax",
			"sub\trsp,38h",
			"mov\trcx,offset __PTestFileName",
			"mov\trdx,40000000h",
			"xor\tr8,r8",
			"xor\tr9,r9",
			"mov\tqword ptr [rsp+20h],2",
			"mov\tqword ptr [rsp+28h],20h",
			"mov\tqword ptr [rsp+30h],0",
			"call\tCreateFile",
			"mov\trsi,rax",
			"mov\trcx,rax",
			"mov\trdx,__pPTestTimesArray",
			"mov\tr8,rbx",
			"shl\tr8,4",
			"mov\tr9,offset __PTestTime",
			"mov\tqword ptr [rsp+20h],0",
			"call\tWriteFile",
			"mov\trcx,rsi",
			"call\tCloseHandle",
			"xor\trcx,rcx",
			"call\tExitProcess"
	};
	private static final String[]	startFunctionText = {
			"mov\t__PTestTime,rax",
			"mov\t__PTestRdx,rdx",
			"rdtsc",
			"shl\trdx,32",
			"or\trax,rdx",
			"xchg\t__PTestTime,rax",
			"mov\trdx,__PTestRdx"
	};
	private static final String[]	stopFunctionText = {
			"mov\t__PTestRax,rax",
			"mov\t__PTestRdx,rdx",
			"rdtsc",
			"shl\trdx,32",
			"or\trax,rdx",
			"sub\trax,__PTestTime",
			"movsd\t__PTestXmm0,xmm0",
			"cvtsi2sd\txmm0,rax",
			"mov\trax,__PTestRax",	// in case i is rax or [rax], etc
			"mov\trdx,__PTestRdx",	// in case i or n is rdx or [rdx], etc
			"mov\trax,index",
			"shl\trax,4",
			"add\trax,__pPTestTimesArray",
			"mov\t__PTestAddress,rax",
			"mov\trax,__PTestRax",	// in case n is rax or [rax], etc
			"mov\trax,n",
			"mov\trdx,__PTestAddress",
			"movsd\t[rdx][8],xmm0",
			"cvtsi2sd\txmm0,rax",
			"movsd\t[rdx],xmm0",
			"mov\trax,__PTestRax",
			"mov\trdx,__PTestRdx",
			"movsd\txmm0,__PTestXmm0"
	};

	public static final PerformanceTestEncoder	instance = new PerformanceTestEncoder();

	private PerformanceTestEncoder() {
		rootFunction = new Function("__PTestEntryPoint", "", ASM.language);
		MultiLineList code = rootFunction.getCode();
		// Remove the empty line
		code.remove(0);
		for (String line : rootFunctionText) {
			ASMLine aLine = ASM.language.createCodeLine(code,code.size(),line.toCharArray());
			code.add(aLine);
			// NOTE: There are no line labels in rootFunction, so they don't need to be added to local scope
		}
		startFunction = new Function("PTestStartTiming", "Starts timing a section of code in a performance test.\nTo ignore the timing, just don't call PTestStopTiming or increment the result index.\nThis can be called again without calling PTestStopTiming in that case.", ASM.language);
		startFunction.getCallingConvention().setValue(Function.INLINE_TEXT);
		code = startFunction.getCode();
		// Remove the empty line
		code.remove(0);
		for (String line : startFunctionText) {
			ASMLine aLine = ASM.language.createCodeLine(code,code.size(),line.toCharArray());
			code.add(aLine);
			// NOTE: There are no line labels in startFunction, so they don't need to be added to local scope
		}
		stopFunction = new Function("PTestStopTiming", "Stops timing a section of code in a performance test\nand saves the time at the specified result index with its\ncorresponding input parameter n.", ASM.language);
		stopFunction.getCallingConvention().setValue(Function.INLINE_TEXT);
		stopFunction.addParameter(new LocalVariable("index","","",stopFunction, CodeFileChunk.LOCATION_PARAMETERS,0));
		stopFunction.addParameter(new LocalVariable("n","","",stopFunction, CodeFileChunk.LOCATION_PARAMETERS,1));
		code = stopFunction.getCode();
		// Remove the empty line
		code.remove(0);
		for (String line : stopFunctionText) {
			ASMLine aLine = ASM.language.createCodeLine(code,code.size(),line.toCharArray());
			code.add(aLine);
			// NOTE: There are no line labels in stopFunction, so they don't need to be added to local scope
		}
		timesArrayVar = new GlobalVariable("__pPTestTimesArray","");
		timesArrayVar.getPieces().add(new GlobalVariable.Piece("ptr qword","?",timesArrayVar,0));
		filenameVar = new GlobalVariable("__PTestFileName","");
		filenameVar.getPieces().add(new GlobalVariable.Piece("byte","\"\",0",filenameVar,0));
		timeVar = new GlobalVariable("__PTestTime","");
		timeVar.getPieces().add(new GlobalVariable.Piece("qword","?",timeVar,0));
		raxVar = new GlobalVariable("__PTestRax","");
		raxVar.getPieces().add(new GlobalVariable.Piece("qword","?",raxVar,0));
		rdxVar = new GlobalVariable("__PTestRdx","");
		rdxVar.getPieces().add(new GlobalVariable.Piece("qword","?",rdxVar,0));
		xmm0Var = new GlobalVariable("__PTestXmm0","");
		xmm0Var.getPieces().add(new GlobalVariable.Piece("qword","?",xmm0Var,0));
		addressVar = new GlobalVariable("__PTestAddress","");
		addressVar.getPieces().add(new GlobalVariable.Piece("qword","?",addressVar,0));
		Scope scope = rootFunction.getScope();
		scope.add(timesArrayVar);
		scope.add(filenameVar);
		scope.add(timeVar);
		scope = startFunction.getScope();
		scope.add(timeVar);
		scope.add(rdxVar);
		scope = stopFunction.getScope();
		scope.add(raxVar);
		scope.add(rdxVar);
		scope.add(timeVar);
		scope.add(xmm0Var);
		scope.add(timesArrayVar);
		scope.add(addressVar);
	}

	/**
	 * Builds and runs a performance test given by the specified Function
	 * @param project the project containing the performance test
	 * @param testFunction the performance test
	 * @return the name of the output file containing the test results, or null if unsuccessful.
	 */
	String runPerformanceTest(Project project,Function testFunction) {
		if (testFunction.getTestFunctionInfo()==null || !(testFunction.getTestFunctionInfo() instanceof PerformanceTestInfo)) {
			ErrorHandler.showMessage("Somehow the function "+testFunction+" tried to be run as a performance test, but it's not one.");
			return null;
		}
		final PerformanceTestInfo testInfo = (PerformanceTestInfo)testFunction.getTestFunctionInfo();
		Section parent = testFunction.getParent();
		// Set up the variables first, because the functions will depend on them having a Language
		timesArrayVar.setParent(parent);
		timesArrayVar.nowHaveGlobalScope();
		String executableName = testFunction.getName().toString();
		String testName = executableName.substring(5);
		final String projectPath = project.getFile().getParent()+File.separatorChar;
		executableName = projectPath+executableName+".exe";
		GregorianCalendar now = new GregorianCalendar();
		int year = now.get(Calendar.YEAR);
		int month = now.get(Calendar.MONTH)+1;
		int day = now.get(Calendar.DAY_OF_MONTH);
		int hour = now.get(Calendar.HOUR_OF_DAY);
		int minute = now.get(Calendar.MINUTE);
		int second = now.get(Calendar.SECOND);
		String dateText = YEAR_FORMAT.format(year)+"-"+NUMBER_FORMAT.format(month)+"-"+NUMBER_FORMAT.format(day)+"_"+NUMBER_FORMAT.format(hour)+"-"+NUMBER_FORMAT.format(minute)+"-"+NUMBER_FORMAT.format(second);
		String scalingText = PerformanceTestInfo.BRIEF_SCALINGS[ArrayUtil.indexOf(PerformanceTestInfo.SCALINGS,testInfo.scaling.toString())];
		final String groupPrefix = ((testInfo.testGroupName.length()>0) ? (testInfo.testGroupName.toString()+"_") : testName+"_");
		final String outputName = ((testInfo.testGroupName.length()>0) ? (testInfo.testGroupName.toString()+"_") : "")+testName+"_"+scalingText+"_"+dateText+".perf";
		filenameVar.setParent(parent);
		ValueLine filenameLine = filenameVar.getPieces().get(0).getInitialValue();
		filenameLine.replace(1,filenameLine.length()-3,outputName);
		filenameVar.nowHaveGlobalScope();
		timeVar.setParent(parent);
		timeVar.nowHaveGlobalScope();
		raxVar.setParent(parent);
		raxVar.nowHaveGlobalScope();
		rdxVar.setParent(parent);
		rdxVar.nowHaveGlobalScope();
		xmm0Var.setParent(parent);
		xmm0Var.nowHaveGlobalScope();
		addressVar.setParent(parent);
		addressVar.nowHaveGlobalScope();
		// Then set up the functions
		rootFunction.setParent(parent);
		Line line = rootFunction.getCode().get(MAX_DATA_LINE_NUM);
		int col = line.toString().indexOf(',');
		line.replace(col+1,line.length(),Long.toString(testInfo.maxNumTestResults.get()));
		line = rootFunction.getCode().get(TEST_FUNCTION_LINE_NUM);
		col = line.toString().indexOf('\t');
		line.replace(col+1,line.length(),testFunction.getName().toString());
		rootFunction.nowHaveGlobalScope();
		startFunction.setParent(parent);
		startFunction.nowHaveGlobalScope();
		stopFunction.setParent(parent);
		stopFunction.nowHaveGlobalScope();
		ArrayList<SyntaxError> errors = new ArrayList<SyntaxError>();
		try {
			long startTime = System.currentTimeMillis();
			Encoder.encodeProject(project,rootFunction,executableName,errors);
			if (!errors.isEmpty()) {
				PwnIDE.editor.showErrors(errors);
			}
			else {
				System.out.println("Compilation Completed in "+(System.currentTimeMillis()-startTime)+"ms with No Errors");
				PwnIDE.editor.runCompiledProgram(executableName,new Thread() {
					public void run() {
						try {
							File outputFile = new File(projectPath+outputName);
							if (outputFile.exists()) {
								File[] files = new File(projectPath).listFiles();
								ArrayList<File> groupFiles = new ArrayList<File>(8);
								for (File file : files) {
									String name = file.getName();
									if (name.startsWith(groupPrefix) && name.endsWith(".perf")) {
										groupFiles.add(file);
									}
								}
								if (!groupFiles.isEmpty()) {
									System.out.println(groupFiles.size()+" data sets");
									double[][] data = new double[groupFiles.size()][];
									for (int i=0;i<groupFiles.size();++i) {
										System.out.println(groupFiles.get(i)+" data:");
										byte[] bytes = FileUtil.readFile(groupFiles.get(i));
										data[i] = ArrayUtil.rawBytesToDoubles(bytes);
										for (int j=0;j<data[i].length;j+=2) {
											System.out.println(data[i][j]+"\t"+data[i][j+1]);
										}
									}
									PerformanceViewer viewer = new PerformanceViewer(groupPrefix.substring(0,groupPrefix.length()-1),groupFiles.toArray(new File[groupFiles.size()]),data);
									viewer.setVisible(true);
								}
								else {
									ErrorHandler.showMessage("Somehow, the test succeeded, but there are no performance test files starting with \""+groupPrefix+"\", sorry.");
								}
							}
							else {
								ErrorHandler.showMessage("The performance test output file "+outputName+"\ncouldn't be read in, sorry.  The test probably crashed.");
							}
						}
						catch (IOException e) {
							ErrorHandler.showMessage("The performance test output file "+outputName+"\ncouldn't be read in, sorry.  The test probably crashed.");
						}
						catch (Throwable t) {
							ErrorHandler.showException(t,"An error occurred when analysing the performance test data\nin "+outputName);
						}
					}
				});
			}
		}
		catch (Throwable t) {
			ErrorHandler.showException(t,"An error occurred when compiling the performance test "+testFunction+".");
			return null;
		}
		finally {
			Scope globalScope = testFunction.getGlobalScope();
			if (globalScope!=null) {
				globalScope.removeObserver(addressVar);
				globalScope.removeObserver(xmm0Var);
				globalScope.removeObserver(rdxVar);
				globalScope.removeObserver(raxVar);
				globalScope.removeObserver(timeVar);
				globalScope.removeObserver(filenameVar);
				globalScope.removeObserver(timesArrayVar);
				globalScope.removeObserver(stopFunction);
				globalScope.removeObserver(startFunction);
				globalScope.removeObserver(rootFunction);
			}
		}
		return outputName;
	}
}
